import numpy as np
import math


def inner_product(u, v, rho):
    N = len(u)
    if len(v) != N:
        raise ValueError("u和v长度不想等")
    s = 0
    for i in range(N):
        s += u[i] * v[i] * rho[i]
    return s


def least_squares_normal_equations(x, y, order=2):
    N = len(x)
    size = order + 1
    G = np.zeros((size, size))
    c = np.zeros((size))
    base = np.zeros((size, N))
    rho = np.ones_like(x)

    # Create the matrix G
    for i in range(size):
        base[i, :] = x**i
    for i in range(size):
        for j in range(size):
            G[i, j] = inner_product(base[i, :], base[j, :], rho)
        c[i] = inner_product(base[i, :], y, rho)

    # Solve the normal equations to obtain coefficients
    coefficients = np.linalg.solve(G, c)

    return coefficients, G


def problemA():
    print()
    print("problemA:")
    print()
    x_data = np.array([0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0,
                       4.5, 5.0, 5.5, 6.0, 6.5, 7.0, 7.5, 8.0, 8.5, 9.0, 9.5, 10.0])  # type: ignore
    y_data = np.array([2.9, 2.7, 4.8, 5.3, 7.1, 7.6, 7.7, 7.6, 9.4,
                       9.0, 9.6, 10.0, 10.2, 9.7, 8.3, 8.4, 9.0, 8.3, 6.6, 6.7, 4.1])

    # x_data = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])  # type: ignore
    # y_data = np.array([256, 201, 159, 61, 77, 40, 17, 25, 103, 156, 222, 345])

    a, G = least_squares_normal_equations(x_data, y_data)
    cond_G = np.linalg.cond(G)
    # Print the coefficients
    print(
        f"The coefficients of the quadratic polynomial are: a0 = {a[0]}, a1 = {a[1]}, a2 = {a[2]}")
    print(f"The condition of G is: {cond_G}")
    return


def gram_schmidt_qr(A):
    m, n = A.shape
    Q = np.zeros((m, n))
    R = np.zeros((n, n))

    for j in range(n):
        v = A[:, j]
        for i in range(j):
            R[i, j] = np.dot(Q[:, i], A[:, j])
            v = v - R[i, j] * Q[:, i]

        R[j, j] = np.linalg.norm(v)
        Q[:, j] = v / R[j, j]

    return Q, R


def least_squares_QR(x, y, order=2):
    N = len(x)
    size = order + 1
    A = np.zeros((N, size))

    # Create the matrix A
    for j in range(size):
        A[:, j] = x**j

    # Q, R = np.linalg.qr(A)
    Q, R = gram_schmidt_qr(A)

    R1 = np.dot(Q.T, A)[0:size, :]
    c = np.dot(Q.T, y)[0:size]

    # Solve the equations to obtain coefficients
    coefficients = np.linalg.solve(R1, c)

    return coefficients, R1


def problemB():
    print()
    print("problemB:")
    print()
    x_data = np.array([0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0,
                       4.5, 5.0, 5.5, 6.0, 6.5, 7.0, 7.5, 8.0, 8.5, 9.0, 9.5, 10.0])  # type: ignore
    y_data = np.array([2.9, 2.7, 4.8, 5.3, 7.1, 7.6, 7.7, 7.6, 9.4,
                       9.0, 9.6, 10.0, 10.2, 9.7, 8.3, 8.4, 9.0, 8.3, 6.6, 6.7, 4.1])

    # x_data = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])  # type: ignore
    # y_data = np.array([256, 201, 159, 61, 77, 40, 17, 25, 103, 156, 222, 345])

    a, R1 = least_squares_QR(x_data, y_data)
    cond_R1 = np.linalg.cond(R1)
    # Print the coefficients
    print(
        f"The coefficients of the quadratic polynomial are: a0 = {a[0]}, a1 = {a[1]}, a2 = {a[2]}")
    print(f"The condition of R1 is: {cond_R1}")
    return


def lu_Gauss_elimination(matrix):
    n = len(matrix)
    L = np.eye(n)
    U = matrix.copy()

    for i in range(n):
        # 检查主元是否为0
        if U[i, i] == 0:
            raise ValueError("主元为零，LU 分解失败")

        # 更新下三角矩阵 L 和上三角矩阵 U
        for j in range(i + 1, n):
            factor = U[j, i] / U[i, i]
            L[j, i] = factor
            U[j, :] -= factor * U[i, :]

    return L, U


def backward_substitution(U, y):
    n = len(y)
    x = np.zeros(n)

    for i in range(n - 1, -1, -1):
        x[i] = y[i] - np.dot(U[i, i + 1:], x[i + 1:])
        x[i] /= U[i, i]

    return x


def forward_substitution(L, b):
    n = len(b)
    x = np.zeros(n)

    for i in range(n):
        x[i] = b[i] - np.dot(L[i, :i], x[:i])
        x[i] /= L[i, i]

    return x


def problemC():
    print()
    print("problemC:")
    print()
    float_format = "{:>" + str(11) + ".6f}"
    np.set_printoptions(linewidth=100)
    # np.set_printoptions(border=False)
    alpha = 1e-10  # Tikhonov parameter
    for k in range(6):
        n = k + 1
        H = np.zeros((n + 1, n + 1))
        R1 = np.zeros((n + 1, 1))
        R2 = np.zeros((n + 1, 1))

        for i in range(n + 1):
            R1[i] = (-1)**i
            for j in range(n + 1):
                if j >= 1 and j <= i:
                    R2[i] += (-1)**(i + j) / j
                H[i, j] = 1 / (i + j + 1)
        L, U = lu_Gauss_elimination(H)

        ln2_machine = math.log(2, math.e)
        t1 = forward_substitution(L, R1)
        beta = backward_substitution(U, t1)
        t2 = forward_substitution(L, R2)
        gamma = backward_substitution(U, t2)
        alpha_m = beta * ln2_machine + gamma
        print(f'n={n} ', end='')
        for num in alpha_m:
            print(float_format.format(num), end=" ")
        print()
        # print(f'n={n}', alpha_m)

        # R_m = R1 * ln2_machine + R2
        # t3 = forward_substitution(L, R_m)
        # alpha_m = backward_substitution(U, t3)
        # print(f'n={n}', alpha_m)

        ln2_e5 = 0.69315
        R_5 = R1 * ln2_e5 + R2
        t4 = forward_substitution(L, R_5)
        alpha_5 = backward_substitution(U, t4)
        print(f'n={n} ', end='')
        for num in alpha_5:
            print(float_format.format(num), end=" ")
        print()
        # print(f'n={n}', alpha_5)

        # alpha_5 = beta * ln2_e5 + gamma
        # print(f'n={n}', alpha_5)

        A = alpha * np.identity(n + 1) + np.dot(H.T, H)
        b = np.dot(H.T, R_5)
        alpha_t = np.linalg.solve(A, b).ravel()
        print(f'n={n} ', end='')
        for num in alpha_t:
            print(float_format.format(num), end=" ")
        print()
        # print(f'n={n}', alpha_t)

        print(f'n={n} ,the condtion number of H is {np.linalg.cond(H)}')
    return


if __name__ == '__main__':
    problemA()
    problemB()
    problemC()
